#!/usr/bin/python

import sys

import numpy as np

from power_iteration_figure import plot_figure
from data import circle_samples, radial_kernel, mutual_knn, rename_clusters

from scipy.sparse import lil_matrix
from scipy.cluster.vq import kmeans2
from scipy.spatial.distance import chebyshev

def get_start_vector(W):
    volume = W.sum()
    v0 = np.array([W.getrow(i).sum() for i in range(W.shape[0])])
    return v0 / volume

def power_iteration(W, vt, tol, max_iter=1000):
    dt = np.zeros(W.shape[0])
    for t in range(max_iter + 1):
        vtlast = vt
        vt = W * vt

        # dt is velocity at t
        dtlast = dt
        dt = vt - vtlast

        # et is acceleration at t maximum norm
        et = chebyshev(dt, dtlast)

        # check convergence
        if et < tol:
            break
    return t, vt

def main(args):
    points = circle_samples()
    W = mutual_knn(points)

    # need the unnormalized version of W
    vt = get_start_vector(W)

    # normalize W
    D = lil_matrix(W.shape)
    D.setdiag([1 / W.getrow(i).sum() for i in range(W.shape[0])])
    W = np.dot(D, W)

    # intermediate vectors for plotting    
    vectors = []
    for t in (1, 5):
        vectors.append(power_iteration(W, vt, 0.0, max_iter=t))

    # run until convergence + cluster
    tol = 1.0e-5 / len(points)
    t, result = power_iteration(W, vt, tol)
    res, idx = kmeans2(result, 3, minit='random')
    idx = rename_clusters(idx)

    vectors.append((t, result))
    plot_figure(points, vectors, idx)

if __name__ == "__main__":
    main(sys.argv)
